#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient, ASCENDING, DESCENDING
import sys

default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    reload(sys)
    sys.setdefaultencoding(default_encoding)


# 连接MongoDB
def InitMongoDB():
    MONGODB_HOST = 'localhost'
    MONGODB_PORT = 27017
    MONGODB_DB_NAME = 'Dictionary'
    MONGODB_COLLECTION_NAME = 'mc3_v_0_2_8'
    connection = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = connection[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return connection, collection


def ParseTSV(MongoDB_Collection, filepath):
    count = 0
    infile = open(filepath, "r")
    headline = infile.readline()
    headline = headline.rstrip("\n")
    splitheadline = headline.split("\t")
    for line in infile:
        line = line.rstrip("\n")
        insert_record = ParseTSVLine(line, splitheadline)
        if len(insert_record) != 0:
            count += 1
            insert_record["ID"] = count
            MongoDB_Collection.insert_one(insert_record)
    infile.close()


def ParseTSVLine(line, splitheadline):
    int_column = set(
        ["Start_Position", "End_Position", "t_depth", "t_ref_count", "t_alt_count", "n_depth", "n_ref_count",
         "n_alt_count", "NCALLERS"])
    comma_column = set(["Existing_variation", "TREMBL", "COSMIC"])
    verticalbar_column = set(["CENTERS", "DBVS"])
    result = {}
    if line != "":
        splitline = line.split("\t")
        for i in xrange(len(splitline)):
            item = splitline[i].strip()
            if item != ".":
                if splitheadline[i] in int_column:
                    result[splitheadline[i]] = int(item)
                elif splitheadline[i] in comma_column:
                    splititems = item.split(",")
                    for splititem in splititems:
                        if result.has_key(splitheadline[i]):
                            result[splitheadline[i]].append(splititem)
                        else:
                            result[splitheadline[i]] = [splititem]
                elif splitheadline[i] in verticalbar_column:
                    splititems = item.split("|")
                    for splititem in splititems:
                        if result.has_key(splitheadline[i]):
                            result[splitheadline[i]].append(splititem)
                        else:
                            result[splitheadline[i]] = [splititem]
                elif splitheadline[i] == "all_effects":
                    if item.endswith(";"):
                        item = item[:-1]
                    splititems = item.split(";")
                    for splititem in splititems:
                        if result.has_key(splitheadline[i]):
                            result[splitheadline[i]].append(splititem)
                        else:
                            result[splitheadline[i]] = [splititem]
                elif splitheadline[i] == "DOMAINS":
                    splititems = item.split(";")
                    for splititem in splititems:
                        item_db = splititem.split(":")[0]
                        item_id = splititem[len(item_db) + 1:]
                        if result.has_key(splitheadline[i]):
                            result[splitheadline[i]].append({"db": item_db, "id": item_id})
                        else:
                            result[splitheadline[i]] = [{"db": item_db, "id": item_id}]
                else:
                    result[splitheadline[i]] = item
    return result


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


if __name__ == "__main__":
    MongoDB_Connection, MongoDB_Collection = InitMongoDB()
    ParseTSV(MongoDB_Collection, "mc3.v0.2.8.PUBLIC.maf")
    CloseMongoDB(MongoDB_Connection)
